# 导入必要的库
from fastapi import FastAPI, UploadFile, File, HTTPException  # FastAPI相关组件
from fastapi.responses import JSONResponse, HTMLResponse  # 响应类型
from fastapi.staticfiles import StaticFiles  # 静态文件服务
from typing import List, Dict, Any  # 类型注解
from paddleocr import PaddleOCR  # PaddleOCR文字识别
from datetime import datetime  # 时间处理
import numpy as np  # 数值计算
import sqlite3  # SQLite数据库
import base64  # Base64编码解码
import cv2  # OpenCV图像处理
import os  # 操作系统接口

# 初始化PaddleOCR，不使用方向分类器，语言设置为中文
ocr = PaddleOCR(use_angle_cls=False, lang="ch")

# 创建FastAPI应用实例
app = FastAPI()

# 挂载静态文件目录，使前端文件可以通过/static路径访问
app.mount("/static", StaticFiles(directory="static"), name="static")

# 设置上传文件保存目录
UPLOAD_DIR = "uploads"
# 确保上传目录存在，如果不存在则创建
os.makedirs(UPLOAD_DIR, exist_ok=True)

# 全局数据库连接变量
db_conn = None

# 应用启动时执行的事件
@app.on_event("startup")
async def startup():
    global db_conn
    # 连接SQLite数据库
    db_conn = sqlite3.connect("paddleocr.sqlite3")
    # 创建存储OCR结果的表（如果不存在）
    db_conn.execute("""
        CREATE TABLE IF NOT EXISTS paddleocr (
            'id' INTEGER NOT NULL UNIQUE,
            'path' TEXT NOT NULL,
            'text' TEXT NOT NULL,
            'json' TEXT,
            PRIMARY KEY('id' AUTOINCREMENT)
        )
    """)

# 应用关闭时执行的事件
@app.on_event("shutdown")
async def shutdown():
    global db_conn
    # 关闭数据库连接
    if db_conn:
        db_conn.close()

# 根路径处理，返回首页HTML
@app.get("/", response_class=HTMLResponse)
async def read_index():
    with open("static/index.html", "r", encoding="utf-8") as f:
        return f.read()

# /main路径处理，返回主页面HTML
@app.get("/main", response_class=HTMLResponse)
async def read_main():
    with open("static/main.html", "r", encoding="utf-8") as f:
        return f.read()

# /main路径处理，返回主页面HTML
@app.get("/three", response_class=HTMLResponse)
async def read_main():
    with open("static/three.html", "r", encoding="utf-8") as f:
        return f.read()

# 将Base64字符串转换为OpenCV图像格式
def base64_to_cv2(base64_str: str) -> np.ndarray:
    try:
        # 处理可能包含MIME类型前缀的Base64字符串
        if "," in base64_str:
            header, base64_data = base64_str.split(",", 1)
            image_data = base64.b64decode(base64_data)
        else:
            image_data = base64.b64decode(base64_str)
        # 将Base64解码后的数据转换为numpy数组
        np_arr = np.frombuffer(image_data, np.uint8)
        # 将numpy数组解码为OpenCV图像格式
        img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        return img
    except Exception as e:
        raise ValueError(f"Base64 解码失败: {str(e)}")

# 处理OCR识别，返回识别结果列表
def ocr_process(image: np.ndarray) -> List[Dict[str, Any]]:
    if image is None:
        raise ValueError("图像数据为空")
    # 检查图像格式是否有效（3维数组且通道数为1、3或4）
    if len(image.shape) != 3 or image.shape[2] not in (1, 3, 4):
        raise ValueError("无效的图像格式")
    # 使用PaddleOCR进行文字识别
    result = ocr.ocr(image)
    ocr_results = []
    # 处理识别结果
    if result and result[0]:
        for line in result:
            if len(line) >= 2:
                ocr_results += line["rec_texts"]
    return ocr_results

# OCR API接口，接收上传的文件
@app.post("/api/ocr/")
async def ocr_api(file: UploadFile = File(...)):
    try:
        # 读取上传的文件内容
        contents = await file.read()
        if not contents:
            raise HTTPException(status_code=400, detail="上传的文件为空")
        
        # 生成保存文件名（使用时间戳）
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        file_extension = os.path.splitext(file.filename)[1]
        saved_filename = f"{timestamp}{file_extension}"
        saved_path = os.path.join(UPLOAD_DIR, saved_filename)
        
        # 将文件保存到上传目录
        with open(saved_path, "wb") as buffer:
            buffer.write(contents)
        
        try:
            # 尝试将文件内容作为Base64字符串解码
            image = base64_to_cv2(contents.decode("utf-8"))
        except UnicodeDecodeError:
            # 如果不是Base64字符串，则直接作为二进制数据解码
            np_arr = np.frombuffer(contents, np.uint8)
            image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
        
        if image is None:
            raise HTTPException(status_code=400, detail="无法解码图像")
        
        # 进行OCR处理
        results = ocr_process(image)
        print(results)
        
        # 将结果存入数据库
        cursor = db_conn.cursor()
        cursor.execute(
            "INSERT INTO paddleocr (text, path) VALUES (?, ?)",
            ("\n".join(results), saved_path)
        )
        db_conn.commit()
        
        return {"status": "success", "results": results}
    except Exception as e:
        # 发生异常时回滚数据库操作
        db_conn.rollback()
        raise HTTPException(status_code=500, detail=str(e))

# 获取所有OCR识别记录的API
@app.get("/paddleocr/")
async def read_items():
    cursor = db_conn.cursor()
    # 查询数据库中的所有记录
    cursor.execute("SELECT * FROM paddleocr")
    return cursor.fetchall()